import math

import requests
import json
import math
from args import *
from tqdm import tqdm
from elasticsearch import Elasticsearch, helpers
import math
client = Elasticsearch(
            hosts=ELASTIC_ADDRESS, basic_auth=(ELASTIC_USERNAME, ELASTIC_PASSWORD),
            timeout=10000)
url = model_url
headers = {
    'Content-Type': 'application/json'
}
from read_data import batch_read

ret = (batch_read("2024-06-01", "2024-08-05"))
print(f'total data: {len(ret)}')
x = []
for r in ret:
    title = r.get("title", "")
    desc = r.get("desc", "")
    x.append(f'{title} {desc}'.strip())
total_times = int(math.ceil(len(ret) / BATCH_SIZE_ES))
vector = []

for i in tqdm(range(total_times)):
    start = i * BATCH_SIZE_ES
    batch_data = x[start:start + BATCH_SIZE_ES]
    batch_data = {
        "data": batch_data,
    }
    json_data = json.dumps(batch_data)
    # 发送 POST 请求
    response = requests.post(url, headers=headers, data=json_data)
    res = response.json()['data']
    vector += res

for i in range(len(x)):
    ret[i]['vector'] = vector[i]
pairs_ = []
for r in ret:
    key = r['key']
    pairs_.append((key, r))
pairs_ = [pairs_]
def update_vectors_in_es(pairs_):
    for pairs in tqdm(pairs_):
        actions = [
            {
                "_op_type": "index",
                "_index": WRITE_NAME,
                "_id": text_id,
                "_source": pair
            }
            for text_id, pair in pairs
        ]
        helpers.bulk(client, actions)

update_vectors_in_es(pairs_)

